import backtrader as bt
import pandas as pd

class MacdStrategy(bt.Strategy):
    """
    一个基于MACD指标并结合基本面过滤的趋势跟踪策略。
    当MACD线上穿信号线（金叉），且最近一期财报的净利润同比增长率为正时，买入。
    当MACD线下穿信号线（死叉）时，卖出。
    """
    params = (
        ('macd1', 12),
        ('macd2', 26),
        ('macdsig', 9),
        ('printlog', True),
        ('financial_data', None), # 接收财务数据的参数
    )

    def __init__(self):
        self.order = None
        self.macd = bt.indicators.MACD(
            self.data.close,
            period_me1=self.p.macd1,
            period_me2=self.p.macd2,
            period_signal=self.p.macdsig
        )
        self.crossover = bt.indicators.CrossOver(self.macd.macd, self.macd.signal)
        
        # 预处理财务数据
        self.financials = None
        if self.p.financial_data is not None and not self.p.financial_data.empty:
            self.financials = self.p.financial_data.copy()
            # 确保公告日期是datetime对象，并将其设为索引
            self.financials['公告日期'] = pd.to_datetime(self.financials['公告日期'])
            self.financials.set_index('公告日期', inplace=True)
            self.financials.sort_index(inplace=True)

    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()}, {txt}')

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f'BUY EXECUTED, Price: {order.executed.price:.2f}')
            else:
                self.log(f'SELL EXECUTED, Price: {order.executed.price:.2f}')
        self.order = None

    def next(self):
        if self.order:
            return

        # 检查基本面是否满足条件
        profit_growth_ok = False
        if self.financials is not None:
            # 获取当前K线日期之前最新的财务数据
            # 'right'表示包含边界，即如果今天有财报，也算进去
            latest_financials = self.financials.truncate(after=self.data.datetime.date(0))
            if not latest_financials.empty:
                # 获取净利润同比增长率(%)
                latest_growth_rate = latest_financials.iloc[-1]['净利润同比增长率(%)']
                if pd.notna(latest_growth_rate) and latest_growth_rate > 0:
                    profit_growth_ok = True
        else:
            # 如果没有财务数据，则默认条件为真，退化为纯技术策略
            profit_growth_ok = True

        if not self.position:
            # 如果出现金叉，并且基本面OK
            if self.crossover > 0 and profit_growth_ok:
                self.log(f'BUY CREATE (Profit Growth OK), Close: {self.data.close[0]:.2f}')
                self.order = self.buy()
            elif self.crossover > 0 and not profit_growth_ok:
                self.log(f'Golden Cross but Profit Growth is NOT OK. No trade.')

        else:
            # 死叉时无条件卖出
            if self.crossover < 0:
                self.log(f'SELL CREATE, Close: {self.data.close[0]:.2f}')
                self.order = self.sell() 